import pickle
import hashlib


class UrlManager:
    def __init__(self):
        self.new_urls = self.load_progress('new_urls.txt')   # 未爬取 url 集合
        self.old_urls = self.load_progress('old_urls.txt')   # 已爬取 url 集合

    def has_new_url(self):
        """
        判断是否有未爬取的 url
        :return: bool
        """
        return self.new_urls_size() != 0

    def get_new_url(self):
        """
        返回一个未爬取的 url
        :return: str
        """
        new_url = self.new_urls.pop()
        m = hashlib.md5()
        m.update(new_url.encode('utf-8'))
        self.old_urls.add(m.hexdigest()[8:-8])

        return new_url

    def add_new_url(self, url):
        """
        添加一个新的 url
        :param url: 单个 url
        :return: None
        """
        if url is None:
            return None
        m = hashlib.md5()
        m.update(url.encode('utf-8'))
        url_md5 = m.hexdigest()[8:-8]
        if (url not in self.new_urls) and (url_md5 not in self.old_urls):
            self.new_urls.add(url)

    def add_new_urls(self, urls):
        """
        添加多个新的url
        :param urls: 多个 url
        :return: None
        """
        if urls is None:
            return None
        for url in urls:
            self.add_new_url(url)

    def new_urls_size(self):
        """
        返回未爬过的 url 集合的大小
        :return: int
        """
        return len(self.new_urls)

    def old_urls_size(self):
        """
        返回已爬过的 url 集合的大小
        :return: int
        """
        return len(self.old_urls)

    def save_progress(self, path, data):
        """
        保存进度
        :param path: 路径
        :return: None
        """
        with open(path, 'wb') as file:
            pickle.dump(data, file)

    def load_progress(self, path):
        """
        从本地文件加载进度
        :param path: 路径
        :return: set
        """
        print('[+] 从文件加载进度{}'.format(path))
        try:
            with open(path, 'rb') as file:
                return pickle.load(file)
        except:
            print('[!] 无进度文件')

        return set()
























